#!/usr/bin/env python3

import os
import re
import shutil
import subprocess
import sys
import time

import common

defaults = {
    'background': False,
    'cpus': 1,
    'wait_gdb': False,
    'debug_vm': None,
    'eval': None,
    'extra_emulator_args': [],
    'gem5_exe_args': '',
    'gem5_script': 'fs',
    'gem5_readfile': '',
    'gem5_restore': None,
    'graphic': False,
    'initramfs': False,
    'initrd': False,
    'kernel_cli': None,
    'kernel_cli_after_dash': None,
    'eval_busybox': None,
    'kgdb': False,
    'kdb': False,
    'kvm': False,
    'memory': '256M',
    'record': False,
    'replay': False,
    'terminal': False,
    'tmux': None,
    'trace': None,
    'userland': None,
    'userland_before': '',
    'vnc': False,
}

def main(args, extra_args=None):
    global defaults
    args = common.resolve_args(defaults, args, extra_args)
    # Common qemu / gem5 logic.
    # nokaslr:
    # * https://unix.stackexchange.com/questions/397939/turning-off-kaslr-to-debug-linux-kernel-using-qemu-and-gdb
    # * https://stackoverflow.com/questions/44612822/unable-to-debug-kernel-with-qemu-gdb/49840927#49840927
    #   Turned on by default since v4.12
    kernel_cli = 'console_msg_format=syslog nokaslr norandmaps panic=-1 printk.devkmsg=on printk.time=y rw'
    if args.kernel_cli is not None:
        kernel_cli += ' {}'.format(args.kernel_cli)
    kernel_cli_after_dash = ''
    extra_emulator_args = []
    extra_qemu_args = []
    if args.debug_vm is not None:
        debug_vm = ['gdb', common.Newline, '-q', common.Newline] + common.shlex_split(args.debug_vm) + ['--args', common.Newline]
    else:
        debug_vm = []
    if args.wait_gdb:
        extra_qemu_args.extend(['-S', common.Newline])
    if args.eval_busybox is not None:
        kernel_cli_after_dash += ' lkmc_eval_base64="{}"'.format(common.base64_encode(args.eval_busybox))
    if args.kernel_cli_after_dash is not None:
        kernel_cli_after_dash += ' {}'.format(args.kernel_cli_after_dash)
    if args.vnc:
        vnc = ['-vnc', ':0', common.Newline]
    else:
        vnc = []
    if args.initrd or args.initramfs:
        ramfs = True
    else:
        ramfs = False
    if args.eval is not None:
        if ramfs:
            initarg = 'rdinit'
        else:
            initarg = 'init'
        kernel_cli += ' {}=/eval_base64.sh'.format(initarg)
        kernel_cli_after_dash += ' lkmc_eval="{}"'.format(common.base64_encode(args.eval))
    if not args.graphic:
        extra_qemu_args.extend(['-nographic', common.Newline])
    console = None
    console_type = None
    console_count = 0
    if args.arch == 'x86_64':
        console_type = 'ttyS'
    elif common.is_arm:
        console_type = 'ttyAMA'
    console = '{}{}'.format(console_type, console_count)
    console_count += 1
    if not (args.arch == 'x86_64' and args.graphic):
        kernel_cli += ' console={}'.format(console)
    extra_console = '{}{}'.format(console_type, console_count)
    console_count += 1
    if args.kdb or args.kgdb:
        kernel_cli += ' kgdbwait'
        if args.kdb:
            if args.graphic:
                kdb_cmd = 'kbd,'
            else:
                kdb_cmd = ''
            kernel_cli += ' kgdboc={}{},115200'.format(kdb_cmd, console)
        if args.kgdb:
            kernel_cli += ' kgdboc={},115200'.format(extra_console)
    if kernel_cli_after_dash:
        kernel_cli += " -{}".format(kernel_cli_after_dash)
    extra_env = {}
    if args.trace is None:
        do_trace = False
        # A dummy value that is already turned on by default and does not produce large output,
        # just to prevent QEMU from emitting a warning that '' is not valid.
        trace_type = 'load_file'
    else:
        do_trace = True
        trace_type = args.trace

    def raise_rootfs_not_found():
        if not args.dry_run:
            raise Exception('Root filesystem not found. Did you build it?\n' \
                            'Tried to use: ' + common.disk_image)
    def raise_image_not_found():
        if not args.dry_run:
            raise Exception('Executable image not found. Did you build it?\n' \
                            'Tried to use: ' + common.image)
    if common.image is None:
        raise Exception('Baremetal ELF file not found. Tried:\n' + '\n'.join(paths))
    cmd = debug_vm.copy()
    if common.emulator == 'gem5':
        if args.baremetal is None:
            if not os.path.exists(common.rootfs_raw_file):
                if not os.path.exists(common.qcow2_file):
                    raise_rootfs_not_found()
                common.raw_to_qcow2(prebuilt=args.prebuilt, reverse=True)
        else:
            if not os.path.exists(common.gem5_fake_iso):
                os.makedirs(os.path.dirname(common.gem5_fake_iso), exist_ok=True)
                common.write_string_to_file(common.gem5_fake_iso, 'a' * 512)
        if not os.path.exists(common.image):
            # This is to run gem5 from a prebuilt download.
            if (not args.baremetal is None) or (not os.path.exists(common.linux_image)):
                raise_image_not_found()
            common.run_cmd([os.path.join(common.extract_vmlinux, common.linux_image)])
        os.makedirs(os.path.dirname(common.gem5_readfile), exist_ok=True)
        common.write_string_to_file(common.gem5_readfile, args.gem5_readfile)
        memory = '{}B'.format(args.memory)
        gem5_exe_args = common.shlex_split(args.gem5_exe_args)
        if do_trace:
            gem5_exe_args.extend(['--debug-flags={}'.format(trace_type), common.Newline])
        extra_env['M5_PATH'] = common.gem5_system_dir
        # https://stackoverflow.com/questions/52312070/how-to-modify-a-file-under-src-python-and-run-it-without-rebuilding-in-gem5/52312071#52312071
        extra_env['M5_OVERRIDE_PY_SOURCE'] = 'true'
        cmd.extend(
            [
                common.executable, common.Newline,
                '--debug-file=trace.txt', common.Newline,
                '--listener-mode', 'on', common.Newline,
                '--outdir', common.m5out_dir, common.Newline,
            ] +
            gem5_exe_args
        )
        if args.userland is not None:
            cmd.extend([
                common.gem5_se_file, common.Newline,
                '-c', common.resolve_userland(args.userland), common.Newline,
            ])
        else:
            if args.gem5_script == 'fs':
                # TODO port
                if args.gem5_restore is not None:
                    cpt_dirs = common.gem_list_checkpoint_dirs()
                    cpt_dir = cpt_dirs[-args.gem5_restore]
                    extra_emulator_args.extend(['-r', str(sorted(cpt_dirs).index(cpt_dir) + 1)])
                cmd.extend([
                    common.gem5_fs_file, common.Newline,
                    '--disk-image', common.disk_image, common.Newline,
                    '--kernel', common.image, common.Newline,
                    '--mem-size', memory, common.Newline,
                    '--num-cpus', str(args.cpus), common.Newline,
                    '--script', common.gem5_readfile, common.Newline,
                ])
                if args.arch == 'x86_64':
                    if args.kvm:
                        cmd.extend(['--cpu-type', 'X86KvmCPU', common.Newline])
                    cmd.extend(['--command-line', 'earlyprintk={} lpj=7999923 root=/dev/sda {}'.format(console, kernel_cli), common.Newline])
                elif common.is_arm:
                    # TODO why is it mandatory to pass mem= here? Not true for QEMU.
                    # Anything smaller than physical blows up as expected, but why can't it auto-detect the right value?
                    cmd.extend([
                        '--command-line', 'earlyprintk=pl011,0x1c090000 lpj=19988480 rw loglevel=8 mem={} root=/dev/sda {}'.format(memory, kernel_cli), common.Newline,
                        '--dtb-filename', os.path.join(common.gem5_system_dir, 'arm', 'dt', 'armv{}_gem5_v1_{}cpu.dtb'.format(common.armv, args.cpus)), common.Newline,
                        '--machine-type', common.machine, common.Newline,
                    ])
                    if args.baremetal is None:
                        cmd.extend([
                            '--param', 'system.panic_on_panic = True', common.Newline])
                    else:
                        cmd.extend(['--bare-metal', common.Newline])
                        if args.arch == 'aarch64':
                            # https://stackoverflow.com/questions/43682311/uart-communication-in-gem5-with-arm-bare-metal/50983650#50983650
                            cmd.extend(['--param', 'system.highest_el_is_64 = True', common.Newline])
                            cmd.extend(['--param', 'system.auto_reset_addr = True', common.Newline])
            elif args.gem5_script == 'biglittle':
                if args.gem5_restore is not None:
                    cpt_dir = common.gem_list_checkpoint_dirs()[-args.gem5_restore]
                    extra_emulator_args.extend(['--restore-from', os.path.join(common.m5out_dir, cpt_dir)])
                cmd.extend([
                    os.path.join(common.gem5_src_dir, 'configs', 'example', 'arm', 'fs_bigLITTLE.py'), common.Newline,
                    '--big-cpus', '2', common.Newline,
                    '--cpu-type', 'atomic', common.Newline,
                    '--disk', common.disk_image, common.Newline,
                    '--dtb', os.path.join(common.gem5_system_dir, 'arm', 'dt', 'armv8_gem5_v1_big_little_2_2.dtb'), common.Newline,
                    '--kernel', common.image, common.Newline,
                    '--little-cpus', '2', common.Newline,
                ])
        if args.wait_gdb:
            # https://stackoverflow.com/questions/49296092/how-to-make-gem5-wait-for-gdb-to-connect-to-reliably-break-at-start-kernel-of-th
            cmd.extend(['--param', 'system.cpu[0].wait_for_remote_gdb = True', common.Newline])
    else:
        qemu_user_and_system_options = [
            '-trace', 'enable={},file={}'.format(trace_type, common.qemu_trace_file), common.Newline,
        ]
        if args.userland is not None:
            if args.wait_gdb:
                debug_args = ['-g', str(common.gdb_port), common.Newline]
            else:
                debug_args = []
            cmd.extend(
                [
                    os.path.join(common.qemu_build_dir, '{}-linux-user'.format(args.arch), 'qemu-{}'.format(args.arch)), common.Newline,
                    '-L', common.target_dir, common.Newline
                ] +
                qemu_user_and_system_options +
                common.shlex_split(args.userland_before) +
                debug_args +
                [
                    common.resolve_userland(args.userland), common.Newline
                ]
            )
        else:
            if not os.path.exists(common.image):
                raise_image_not_found()
            extra_emulator_args.extend(extra_qemu_args)
            common.make_run_dirs()
            if args.prebuilt:
                qemu_executable = common.qemu_executable_basename
                qemu_found = shutil.which(qemu_executable) is not None
            else:
                qemu_executable = common.qemu_executable
                qemu_found = os.path.exists(qemu_executable)
            if not qemu_found and not args.dry_run:
                raise Exception('QEMU executable not found, did you forget to build or install it?\n' \
                                'Tried to use: ' + qemu_executable)
            if args.debug_vm:
                serial_monitor = []
            else:
                if args.background:
                    serial_monitor = ['-serial', 'file:{}'.format(common.qemu_background_serial_file), common.Newline]
                else:
                    serial_monitor = ['-serial', 'mon:stdio', common.Newline]
            if args.kvm:
                extra_emulator_args.extend(['-enable-kvm', common.Newline])
            extra_emulator_args.extend(['-serial', 'tcp::{},server,nowait'.format(common.extra_serial_port), common.Newline])
            virtfs_data = [
                (common.p9_dir, 'host_data'),
                (common.out_dir, 'host_out'),
                (common.out_rootfs_overlay_dir, 'host_out_rootfs_overlay'),
                (common.rootfs_overlay_dir, 'host_rootfs_overlay'),
            ]
            virtfs_cmd = []
            for virtfs_dir, virtfs_tag in virtfs_data:
                if os.path.exists(virtfs_dir):
                    virtfs_cmd.extend([
                        '-virtfs',
                        'local,path={virtfs_dir},mount_tag={virtfs_tag},security_model=mapped,id={virtfs_tag}' \
                            .format(virtfs_dir=virtfs_dir, virtfs_tag=virtfs_tag),
                        common.Newline,
                    ])
            cmd.extend(
                [
                    qemu_executable, common.Newline,
                    '-device', 'rtl8139,netdev=net0', common.Newline,
                    '-gdb', 'tcp::{}'.format(common.gdb_port), common.Newline,
                    '-kernel', common.image, common.Newline,
                    '-m', args.memory, common.Newline,
                    '-monitor', 'telnet::{},server,nowait'.format(common.qemu_monitor_port), common.Newline,
                    '-netdev', 'user,hostfwd=tcp::{}-:{},hostfwd=tcp::{}-:22,id=net0'.format(common.qemu_hostfwd_generic_port, common.qemu_hostfwd_generic_port, common.qemu_hostfwd_ssh_port), common.Newline,
                    '-no-reboot', common.Newline,
                    '-smp', str(args.cpus), common.Newline,
                ] +
                virtfs_cmd +
                qemu_user_and_system_options +
                serial_monitor +
                vnc
            )
            if args.initrd:
                extra_emulator_args.extend(['-initrd', os.path.join(common.buildroot_images_dir, 'rootfs.cpio')])
            rr = args.record or args.replay
            if ramfs:
                # TODO why is this needed, and why any string works.
                root = 'root=/dev/anything'
            else:
                if rr:
                    driveif = 'none'
                    rrid = ',id=img-direct'
                    root = 'root=/dev/sda'
                    snapshot = ''
                else:
                    driveif = 'virtio'
                    root = 'root=/dev/vda'
                    rrid = ''
                    snapshot = ',snapshot'
                if args.baremetal is None:
                    if not os.path.exists(common.qcow2_file):
                        if not os.path.exists(common.rootfs_raw_file):
                            raise_rootfs_not_found()
                        common.raw_to_qcow2(prebuilt=args.prebuilt)
                    extra_emulator_args.extend([
                        '-drive',
                        'file={},format=qcow2,if={}{}{}'.format(common.disk_image, driveif, snapshot, rrid),
                        common.Newline,
                    ])
                    if rr:
                        extra_emulator_args.extend([
                        '-drive', 'driver=blkreplay,if=none,image=img-direct,id=img-blkreplay', common.Newline,
                        '-device', 'ide-hd,drive=img-blkreplay', common.Newline,
                    ])
            if rr:
                extra_emulator_args.extend([
                    '-object', 'filter-replay,id=replay,netdev=net0',
                    '-icount', 'shift=7,rr={},rrfile={}'.format('record' if args.record else 'replay', common.qemu_rrfile),
                ])
                virtio_gpu_pci = []
            else:
                virtio_gpu_pci = ['-device', 'virtio-gpu-pci', common.Newline]
            if args.arch == 'x86_64':
                append = ['-append', '{} nopat {}'.format(root, kernel_cli), common.Newline]
                cmd.extend([
                    '-M', common.machine, common.Newline,
                    '-device', 'edu', common.Newline,
                ])
            elif common.is_arm:
                extra_emulator_args.extend(['-semihosting', common.Newline])
                if args.arch == 'arm':
                    cpu = 'cortex-a15'
                else:
                    cpu = 'cortex-a57'
                append = ['-append', '{} {}'.format(root, kernel_cli), common.Newline]
                cmd.extend(
                    [
                        # highmem=off needed since v3.0.0 due to:
                        # http://lists.nongnu.org/archive/html/qemu-discuss/2018-08/msg00034.html
                        '-M', '{},highmem=off'.format(common.machine), common.Newline,
                        '-cpu', cpu, common.Newline,
                    ] +
                    virtio_gpu_pci
                )
            if args.baremetal is None:
                cmd.extend(append)
    if args.tmux is not None:
        if common.emulator == 'gem5':
            subprocess.Popen([os.path.join(common.root_dir, 'tmu'),
                'sleep 2;./gem5-shell -n {} {}' \
                .format(args.run_id, args.tmux)
            ])
        elif args.wait_gdb:
            # TODO find a nicer way to forward all those args automatically.
            # Part of me wants to: https://github.com/jonathanslenders/pymux
            # but it cannot be used as a library properly it seems, and it is
            # slower than tmux.
            subprocess.Popen([os.path.join(common.root_dir, 'tmu'),
                "sleep 2;./run-gdb --arch '{}' --linux-build-id '{}' --run-id '{}' {}" \
                .format(
                    args.arch,
                    args.linux_build_id,
                    args.run_id,
                    args.tmux
                )
            ])
    cmd.extend(extra_emulator_args)
    cmd.extend(args.extra_emulator_args)
    if debug_vm or args.terminal:
        out_file = None
    else:
        out_file = common.termout_file
    common.run_cmd(cmd, cmd_file=common.run_cmd_file, out_file=out_file, extra_env=extra_env)
    # Check if guest panicked.
    if common.emulator == 'gem5':
        # We have to do some parsing here because gem5 exits with status 0 even when panic happens.
        # Grepping for '^panic: ' does not work because some errors don't show that message.
        panic_msg = b'--- BEGIN LIBC BACKTRACE ---$'
    else:
        panic_msg = b'Kernel panic - not syncing'
    panic_re = re.compile(panic_msg)
    error_string_found = False
    if out_file is not None and not args.dry_run:
        with open(common.termout_file, 'br') as logfile:
            for line in logfile:
                if panic_re.search(line):
                    error_string_found = True
        if os.path.exists(common.guest_terminal_file):
            with open(common.guest_terminal_file, 'br') as logfile:
                lines = logfile.readlines()
                if lines:
                    last_line = lines[-1]
                    if last_line.rstrip() == common.magic_fail_string:
                        error_string_found = True
        if error_string_found:
            common.log_error('simulation error detected by parsing logs')
            return 1
    return 0

def get_argparse():
    parser = common.get_argparse(argparse_args={'description':'Run Linux on an emulator'})
    init_group = parser.add_mutually_exclusive_group()
    kvm_group = parser.add_mutually_exclusive_group()
    parser.add_argument(
        '--background', default=defaults['background'], action='store_true',
        help='''\
Send QEMU output to a file instead of the terminal so it does not require a
terminal attached to run on the background. Interactive input cannot be given.
TODO: use a port instead. If only there was a way to redirect a serial to multiple
places, both to a port and a file? We use the file currently to be able to have
any output at all.
https://superuser.com/questions/1373226/how-to-redirect-qemu-serial-output-to-both-a-file-and-the-terminal-or-a-port
'''
    )
    parser.add_argument(
        '-c', '--cpus', default=defaults['cpus'], type=int,
        help='Number of guest CPUs to emulate. Default: %(default)s'
    )
    parser.add_argument(
        '-D', '--debug-vm', default=defaults['debug_vm'], nargs='?', action='store', const='',
        help='Run GDB on the emulator itself.'
    )
    parser.add_argument(
        '-E', '--eval',
        help='''\
Replace the normal init with a minimal init that just evals the given string.
See: https://github.com/cirosantilli/linux-kernel-module-cheat#replace-init
'''
    )
    parser.add_argument(
        '-e', '--kernel-cli',
        help='''\
Pass an extra Linux kernel command line options, and place them before
the dash separator `-`. Only options that come before the `-`, i.e.
"standard" options, should be passed with this option.
Example: `./run -a arm -e 'init=/poweroff.out'`
'''
    )
    parser.add_argument(
        '-F', '--eval-after-init',
        help='''\
Pass a base64 encoded command line parameter that gets evalled by the Busybox init.
See: https://github.com/cirosantilli/linux-kernel-module-cheat#init-busybox
'''
    )
    parser.add_argument(
        '-f', '--kernel-cli-after-dash',
        help='''\
Pass an extra Linux kernel command line options, add a dash `-`
separator, and place the options after the dash. Intended for custom
options understood by our `init` scripts, most of which are prefixed
by `lkmc_`.
Example: `./run --kernel-cli-after-dash 'lkmc_eval="wget google.com" lkmc_lala=y'`
Mnenomic: `-f` comes after `-e`.
'''
    )
    parser.add_argument(
        '-G', '--gem5-exe-args', default=defaults['gem5_exe_args'],
        help='''\
Pass extra options to the gem5 executable.
Do not confuse with the arguments passed to config scripts,
like `fs.py`. Example:
./run -G '--debug-flags=Exec --debug' --gem5 -- --cpu-type=HPI --caches
will run:
gem.op5 --debug-flags=Exec fs.py --cpu-type=HPI --caches
'''
    )
    parser.add_argument(
        '--gem5-script', default=defaults['gem5_script'], choices=['fs', 'biglittle'],
        help='Which gem5 script to use'
    )
    parser.add_argument(
        '--gem5-readfile', default=defaults['gem5_readfile'],
        help='Set the contents of m5 readfile to this string.'
    )
    init_group.add_argument(
        '-I', '--initramfs', default=defaults['initramfs'], action='store_true',
        help='Use initramfs instead of a root filesystem'
    )
    init_group.add_argument(
        '-i', '--initrd', default=defaults['initrd'], action='store_true',
        help='Use initrd instead of a root filesystem'
    )
    kvm_group.add_argument(
        '-K', '--kvm', default=defaults['kvm'], action='store_true',
        help='Use KVM. Only works if guest arch == host arch'
    )
    parser.add_argument(
        '--kgdb', default=defaults['kgdb'], action='store_true'
    )
    parser.add_argument(
        '--kdb', default=defaults['kdb'], action='store_true'
    )
    parser.add_argument(
        '-l', '--gem5-restore', type=int,
        help='''\
Restore the nth most recently taken gem5 checkpoint according to directory
timestamps.
'''
    )
    parser.add_argument(
        '-m', '--memory', default=defaults['memory'],
        help='''\
Set the memory size of the guest. E.g.: `-m 512M`. We try to keep the default
at the minimal ammount amount that boots all archs. Anything lower could lead
some arch to fail to boot.
Default: %(default)s
'''
    )
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        '-R', '--replay', default=defaults['replay'], action='store_true',
        help='Replay a QEMU run record deterministically'
    )
    group.add_argument(
        '-r', '--record', default=defaults['record'], action='store_true',
        help='Record a QEMU run record for later replay with `-R`'
    )
    parser.add_argument(
        '-T', '--trace',
        help='''\
Set trace events to be enabled. If not given, gem5 tracing is completely
disabled, while QEMU tracing is enabled but uses default traces that are very
rare and don't affect performance, because `./configure
--enable-trace-backends=simple` seems to enable some traces by default, e.g.
`pr_manager_run`, and I don't know how to get rid of them.
'''
    )
    init_group.add_argument(
        '--terminal', default=defaults['terminal'], action='store_true',
        help='''\
Output to the terminal, don't pipe to tee as the default.
Does not save the output to a file, but allows you to use debuggers.
Set automatically by --debug-vm, but you still need this option to debug
gem5 Python scripts.
'''
    )
    parser.add_argument(
        '-t', '--tmux', default=defaults['tmux'], nargs='?', action='store', const='',
        help='''\
Create a tmux split the window. You must already be inside of a `tmux` session
to use this option:
* on the main window, run the emulator as usual
* on the split:
** if on QEMU and `-d` is given, GDB
** if on gem5, the gem5 terminal
If values are given to this option, pass those as parameters
to the program running on the split.
'''
    )
    parser.add_argument(
        '-u', '--userland', default=defaults['userland'],
        help='''\
Run the given userland executable in user mode instead of booting the Linux kernel
in full system mode. In gem5, user mode is called Syscall Emulation (SE) mode and
uses se.py.

Path resolution is similar to --baremetal.
'''
    )
    parser.add_argument(
        '--userland-before', default=defaults['userland_before'],
        help='''\
Pass these arguments to the QEMU user mode CLI before the program to execute.
This is required with --userland since arguments that come at the end are interpreted
as command line arguments to that executable.
'''
    )
    kvm_group.add_argument(
        '-w', '--wait-gdb', default=defaults['wait_gdb'], action='store_true',
        help='Wait for GDB to connect before starting execution'
    )
    parser.add_argument(
        '-x', '--graphic', default=defaults['graphic'], action='store_true',
        help='Run in graphic mode. Mnemonic: X11'
    )
    parser.add_argument(
        '-V', '--vnc', default=defaults['vnc'], action='store_true',
        help='''\
Run QEMU with VNC instead of the default SDL. Connect to it with:
`vinagre localhost:5900`.
'''
    )
    parser.add_argument(
        'extra_emulator_args', nargs='*', default=defaults['extra_emulator_args'],
        help='Extra options to append at the end of the emulator command line'
    )
    return parser

if __name__ == '__main__':
    parser = get_argparse()
    args = common.setup(parser)
    start_time = time.time()
    exit_status = main(args)
    end_time = time.time()
    common.print_time(end_time - start_time)
    sys.exit(exit_status)
